import os, sys
from typing import List, Optional
import datetime
import json
import random
import logging
import numpy as np
import toml

from pydantic import BaseModel
from fastapi import FastAPI, APIRouter
from fastapi.middleware.cors import CORSMiddleware

from snowflake import generator
from engine import GptEngine

def setup_logger(save_dir, name=None, log_level=logging.DEBUG):
    if not name:
        logger = logging.getLogger()
        name = 'api_svr'
    else:
        logger = logging.getLogger(name=name)
    logger.setLevel(log_level)

    ch = logging.StreamHandler(stream=sys.stdout)
    ch.setLevel(log_level)
    formatter = logging.Formatter("%(asctime)s %(name)s %(levelname)s: %(message)s")
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    if save_dir:
        if not os.path.exists(save_dir):
            os.makedirs(save_dir)
        fh = logging.FileHandler(os.path.join(save_dir, "api_log.txt"), mode='a+')
        fh.setLevel(log_level)
        fh.setFormatter(formatter)
        logger.addHandler(fh)

    return logger


config_path = './config.toml'
conf = toml.load(config_path)
setup_logger(conf['main']['log_dir'], log_level=conf['main']['log_level'])
logging.info('server config: %s', conf)

model_conf = conf['model']
engine = GptEngine(
        model_name_or_path=model_conf['model_name_or_path'],
        tokenizer_path     = None,
        device             = model_conf['device'],
        fp16               = model_conf['fp16'],
        seed               = model_conf['seed'])

id_gen = generator(1, 1)
app = FastAPI()

origins = [
    "*",
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# api schema
class ReqParam(BaseModel):
    prompt                  : str
    req_id                  : Optional[str] = None
    max_length              : Optional[int] = 0
    min_length              : Optional[int] = 0
    do_sample               : Optional[bool] = True
    top_p                   : Optional[float] = 0.
    top_k                   : Optional[int] = 0
    temperature             : Optional[float] = 0
    stop_words              : Optional[List[str]] = None
    length_penalty          : Optional[float] = None
    repetition_penalty      : Optional[float] = None
    num_return_sequences    : Optional[int] = 1

class Response(BaseModel):
    code        : int
    log_id      : str
    err_msg     : Optional[str] = ''
    texts       : List[str] = None

api_router = APIRouter()
@api_router.post("/generate", response_model=Response)
async def generate(req_param: ReqParam):
    log_id = req_param.req_id or f'{next(id_gen):x}'
    if not req_param.prompt:
        err_msg = 'empty prompt, please check it'
        logging.warn('%s, %s', log_id, err_msg)
        return Response(code=-1, log_id=log_id, err_msg=err_msg)

    try:
        answers = engine.generate(
                req_param.prompt,
                top_k = req_param.top_k or model_conf['top_k'],
                top_p = req_param.top_p or model_conf['top_p'],
                max_len = req_param.max_length or model_conf['max_len'],
                min_len = req_param.min_length or model_conf['min_len'],
                temperature = req_param.temperature or model_conf['temperature'],
                length_penalty = req_param.length_penalty or \
                        model_conf['length_penalty'],
                repetition_penalty = req_param.repetition_penalty or \
                        model_conf['repetition_penalty'],
                do_sample = req_param.do_sample,
                num_return_sequences = req_param.num_return_sequences or \
                        model_conf['num_return_sequences'],
                stop_words = req_param.stop_words or model_conf['stop_words']
        )
        rsp = Response(code=0, log_id=log_id, texts=answers)
        logging.info('(%s): %s <-> %s', log_id, req_param, rsp)
        return rsp

    except Exception as err:
        logging.exception('%s, %s', log_id, err)
        return Response(code=-1, log_id=log_id, err_msg='system error')

app.include_router(api_router, prefix='/api/v0')

if __name__ == '__main__':
    import uvicorn
    uvicorn.run(app="api:app", port=conf['main']['port'], reload=False, debug=True)
